#include "stdafx.h"
#include "Scene.h"
#include "RayTraceSystem.h"
#include <ppl.h>
namespace RayTracePro
{
	using namespace Concurrency;
	Scene::Scene(RayTraceSystem * sys)
	{
		this->system = sys;
		dataInGpu = false;
		initBvhTracingKernel = GpuDevice::CreateKernel("InitBvhTracing");
		traverseBvhKernel = GpuDevice::CreateKernel("TraverseBVH");
		computeHeadFlagsKernel = GpuDevice::CreateKernel("ComputeHeadFlags");
		compressGroups1Kernel = GpuDevice::CreateKernel("CompressGroups_FirstStep");
		compressGroups2Kernel = GpuDevice::CreateKernel("CompressGroups_SecondStep");
	}

	void Scene::LoadMesh(const String & fileName)
	{
		GeometryLoader::Init();
		mesh = new MeshGeometry();
		GeometryLoader::LoadGeometryFromFile(system, *mesh, fileName);
		GeometryLoader::Finalize();
		auto kdFileName = fileName + L".kd";
		if (FileExists(kdFileName))
			mesh->LoadKdTree(kdFileName);
		else
			mesh->UpdateAccelerateStructure();
		mesh->SaveKdTree(kdFileName);
	}

	struct int4
	{
		int x,y,z,w;
	};

	void TraverseBvh(const Bvh & bvh, Ray ray, TraverseResult& continuation)
	{
		BvhNode curNode;
		unsigned int curNodeId = continuation.nodeId;
		unsigned int curParentId = BvhNode::Nil;
		unsigned char (&stack)[BVH_MAX_CHILDREN] = continuation.Stack;
		int stackPtr = (int)continuation.stackPtr;
		continuation.LeafId = -1;
		continuation.MeshId = -1;
		while (curNodeId != BvhNode::Nil && (curNode = bvh.Nodes[curNodeId]).Flags != 1)
		{
			float minT = FLT_MAX;
			int minId = -1;
			unsigned char &childFlags = stack[stackPtr];
			for (int i = 0; i<BVH_MAX_CHILDREN; i++)
			{
				if (curNode.Children[i] == -1)
					break;
				if ((childFlags & (1<<i)) == 0)
				{
					float tMin, tMax;
					if (RayBBoxIntersection(bvh.Nodes[curNode.Children[i]].Bounds, ray, tMin, tMax) && tMin < continuation.tMin)
					{
						if (tMin < minT)
						{
							minT = tMin;
							minId = i;
						}
					}
					else
					{
						childFlags |= (1<<i);
					}
				}
			}
			if (minId != -1)
			{
				curParentId = curNodeId;
				stackPtr ++;
				childFlags |= (1<<minId);
				curNodeId = curNode.Children[minId];
			}
			else
			{
				childFlags = 0;
				stackPtr--;
				curNodeId = curNode.Parent;
				curParentId = curNodeId;
			}
		}
		if (curNodeId != BvhNode::Nil)
		{
			continuation.LeafId = curNode.Children[0];
			continuation.MeshId = curNode.Children[1];
			stackPtr--;
		}
		continuation.nodeId = curParentId;
		continuation.stackPtr = stackPtr&255;
	}

	void InitContinuation(int size, DeviceList<TraverseResult> & continuation)
	{
		for (int i = 0; i<size; i++)
		{
			continuation[i].MeshId = -1;
			continuation[i].LeafId = -1;
			continuation[i].nodeId = 0;
			continuation[i].stackPtr = 0;
			continuation[i].tMin = FLT_MAX;
			continuation[i].Stack0 = 0;
			continuation[i].Stack1 = 0;
		}
	}

	void ComputeHeadFlags(int size, DeviceList<int> & headFlags, DeviceList<int> & scanHeadFlags, DeviceList<TraverseResult> & continuation)
	{
		parallel_for (0, size, [&](int i)
		{
			if (i == 0 || continuation[i-1].MeshId != continuation[i].MeshId)
				headFlags[i] = 1;
			else
			{
				headFlags[i] = 0;
			}
			scanHeadFlags[i] = headFlags[i];
		});
	}

	void Scene::ComputeHeadFlags_GPU(int size, DeviceList<int> & headFlags, DeviceList<TraverseResult> & continuation)
	{
		GpuDevice::InvokeKernel(computeHeadFlagsKernel, Align(size, 256), 256)
			(size, continuation.DeviceBuffer(false), headFlags.DeviceBuffer(false));
	}

	void CompressGroups(int size, DeviceList<int> & headFlags, DeviceList<int> & scanHeadFlags, DeviceList<SortGroup> & groups, DeviceList<TraverseResult> & continuation)
	{
		parallel_for(0, size, [&](int i)
		{
			if (headFlags[i] == 1)
			{
				groups[scanHeadFlags[i]].Base = i;
				groups[scanHeadFlags[i]].GroupId = continuation[i].MeshId;
			}
		});
		parallel_for(0, groups.Count(), [&](int i)
		{
			if (i<groups.Count()-1)
				groups[i].Count = groups[i+1].Base-groups[i].Base;
			else
			{
				groups[i].Count = size-groups[i].Base;
			}
		});
	}

	void Scene::CompressGroups_GPU(int size, DeviceList<int> & headFlags, DeviceList<int> & scanHeadFlags, DeviceList<SortGroup> & groups, DeviceList<TraverseResult> & continuation)
	{
		GpuDevice::InvokeKernel(compressGroups1Kernel, Align(size,256), 256)
			(size, headFlags.DeviceBuffer(false), scanHeadFlags.DeviceBuffer(false), groups.DeviceBuffer(false), continuation.DeviceBuffer(false));
		GpuDevice::InvokeKernel(compressGroups2Kernel, Align(groups.Count(),128), 128)
			(groups.Count(), size, groups.DeviceBuffer(false));
	}

	void DecompressGroups(int size, DeviceList<int> & headFlags, DeviceList<int> & scanHeadFlags, DeviceList<SortGroup> &groups, DeviceList<int> & skeleton)
	{
		for (int i = 0; i<size; i++)
		{
			headFlags[i] = 0;
			scanHeadFlags[i] = 0;
		}
		for (int i = 0; i<groups.Count(); i++)
		{
			headFlags[groups[i].ScannedCount] = 1;
			scanHeadFlags[groups[i].ScannedCount] = 1;
		}
		ExclusiveScan(scanHeadFlags);
		parallel_for (0, size, [&](int i)
		{
			if (headFlags[i])
				skeleton[i] = groups[scanHeadFlags[i]].Base;
			else
			{
				skeleton[i] = 1;
			}
		});
		// segscan(skeleton, headFlags)
		for (int i = 0; i<skeleton.Count(); i++)
		{
			if (!headFlags[i])
				skeleton[i] += skeleton[i-1];
		}

	}

	void Scene::TraceRay_BVH(DeviceList<DifferentialGeometry> & result, 
							 DeviceList<Ray> & rays, 
						 	 DeviceList<RayDifferential> & rayDiffs,
							 DeviceList<TraverseResult> & continuation,
							 DeviceList<int> & headFlags, 
							 DeviceList<int> & scanHeadFlags,
							 DeviceList<SortGroup> & groups,
							 DeviceList<int> & skeleton)
	{
		bool done = true;
		memset(result.Buffer(), 0, sizeof(DifferentialGeometry)*result.Count());
		for (int i = 0; i<result.Count(); i++)
		{
			result[i].MaterialId = -1;
		}
		skeleton.SetSize(rays.Count());
		headFlags.SetSize(rays.Count());
		scanHeadFlags.SetSize(rays.Count());
		continuation.SetSize(rays.Count());
		InitContinuation(rays.Count(), continuation);
		do
		{
			system->StartBenchmark("TraverseBVH");
			done = true;
			RayTracePro::Bvh & bvh = *this->Bvh;
			parallel_for(0, rays.Count(), [&](int i)
			{
				TraverseBvh(bvh, rays[i], continuation[i]);
			});
			system->StopBenchmark();
			system->StartBenchmark("SortRays");
			system->StartBenchmark("HeadFlags+Compress");
			ComputeHeadFlags(rays.Count(), headFlags, scanHeadFlags, continuation);
			// scan
			int groupCount = ExclusiveScan(scanHeadFlags);
			groups.SetSize(groupCount);
			CompressGroups(rays.Count(), headFlags, scanHeadFlags, groups, continuation);
			system->StopBenchmark();
			system->StartBenchmark("Sorting");
			// sort by groupName
			parallel_radixsort(groups.Buffer(), groups.Buffer()+groups.Count(), [](const SortGroup &g){return g.GroupId;});
			//std::sort(groups.Buffer(), groups.Buffer()+groups.Count(), [](const SortGroup &g1, const SortGroup & g2) {return g1.GroupId<g2.GroupId;});
			system->StopBenchmark();
			// build skeleton
			int ptr = 0;
			for (int i = 0; i<groups.Count(); i++)
			{
				int base = groups[i].Base;
				for (int j = 0; j<groups[i].Count; j++)
				{
					skeleton[ptr++] = base + j;
				}
			}
			system->StopBenchmark();
			/*
				// exclusive scan: 
				groups[0].ScannedCount = groups[0].Count;
				for (int i = 1; i<groups.Count(); i++)
				{
					groups[i].ScannedCount = groups[i].Count + groups[i-1].ScannedCount;
				}
				for (int i = groups.Count()-1; i>0; i--)
				{
					groups[i].ScannedCount = groups[i-1].ScannedCount;
				}
				groups[0].ScannedCount = 0;
				// read back groups async
				// ...
				// decompress
				system->StartBenchmark("Decompress");
				DecompressGroups(rays.Count(), headFlags, scanHeadFlags, groups, skeleton);
				system->StopBenchmark();
			*/
			// trace mesh
			system->StartBenchmark("TraceMesh");
			int gPtr = 0;
			int start = 0;
			while (gPtr < groups.Count())
			{
				int end = start + groups[gPtr].Count;
				while (gPtr + 1 < groups.Count() && groups[gPtr+1].GroupId == groups[gPtr].GroupId)
				{
					gPtr++;
					end += groups[gPtr].Count;
				}
				// MeshTrace skeleton[start, end]
				if (groups[gPtr].GroupId != -1)
				{
					done = false;
					RayTraceSystem * sys = system;				
					parallel_for (start, end, [&](int i)
					{
						DifferentialGeometry inter;
						float t;
						int geomId = continuation[skeleton[i]].LeafId;
						// transform ray
						Ray ray = rays[skeleton[i]];
						ray = TransformRay(Instances[geomId].TransformInverse, ray);
						if (meshes[groups[gPtr].GroupId]->ComputeIntersection(sys, 0, inter, ray, rayDiffs[skeleton[i]], t))
						{
							if (t < continuation[skeleton[i]].tMin)
							{
								continuation[skeleton[i]].tMin = t;
								inter = TransformIntersection(Instances[geomId].Transform, inter);
								result[skeleton[i]] = inter;
							}
						}
					});
				}
				start = end;
				gPtr++;
			}
			system->StopBenchmark();
		} while (!done);
	}

	void Scene::InitBvhTracing_GPU(int size, DeviceList<TraverseResult> & continuation, DeviceList<DifferentialGeometry> & result)
	{
		GpuDevice::InvokeKernel(initBvhTracingKernel, Align(size,256), 256)
			(size, continuation.DeviceBuffer(false), result.DeviceBuffer(false));
	}

	void Scene::TraceRay_BVH_GPU(DeviceList<DifferentialGeometry> & result, 
								 DeviceList<Ray> & rays, 
								 DeviceList<RayDifferential> & rayDiffs,
								 DeviceList<TraverseResult> & continuation,
								 DeviceList<int> & headFlags, 
								 DeviceList<int> & scanHeadFlags,
								 DeviceList<SortGroup> & groups,
								 DeviceList<int> & skeleton,
								 DeviceList<Vec4> & intermediateIntersections)
	{
		bool done = true;
		skeleton.SetSize(rays.Count());
		GpuDevice::AllocScanList(headFlags, rays.Count());
		scanHeadFlags.SetSize(headFlags.Count());
		continuation.SetSize(rays.Count());
		if (!dataInGpu)
		{
			Bvh->Nodes.Refresh(true, false);
			Instances.Refresh(true, false);
			dataInGpu = true;
		}
		InitBvhTracing_GPU(rays.Count(), continuation, result);
		do
		{
			system->StartBenchmark("TraverseBVH");
			done = true;
			GpuDevice::InvokeKernel(traverseBvhKernel, Align(rays.Count(), 128), 128)
				(rays.Count(), system->GeometryBias, rays.DeviceBuffer(false), Bvh->Nodes.DeviceBuffer(false), continuation.DeviceBuffer(false));
			system->StopBenchmark();
			system->StartBenchmark("SortRays");
			ComputeHeadFlags_GPU(rays.Count(), headFlags, continuation);
			// scan
			int headFlagLast = 0, scanHeadFlagLast = 0;
			GpuDevice::GetQueue().enqueueReadBuffer(headFlags.DeviceBuffer(false), false, (headFlags.Count()-1)*sizeof(int), sizeof(int), &headFlagLast);
			GpuDevice::ExclusiveScan(scanHeadFlags.DeviceBuffer(false), headFlags.DeviceBuffer(false), scanHeadFlags.Count());
			GpuDevice::GetQueue().enqueueReadBuffer(scanHeadFlags.DeviceBuffer(false), true, (headFlags.Count()-1)*sizeof(int), sizeof(int), &scanHeadFlagLast);
			int groupCount = headFlagLast + scanHeadFlagLast;
			groups.SetSize(groupCount);
			CompressGroups_GPU(rays.Count(), headFlags, scanHeadFlags, groups, continuation);
			groups.Synchronize(true);
			// sort by groupName
			parallel_radixsort(groups.Buffer(), groups.Buffer()+groups.Count(), [](const SortGroup &g){return g.GroupId;});
			// build skeleton
			int ptr = 0;
			for (int i = 0; i < groups.Count(); i++)
			{
				int base = groups[i].Base;
				for (int j = 0; j<groups[i].Count; j++)
				{
					skeleton[ptr++] = base + j;
				}
			}
			skeleton.Refresh(true, false);
			// trace mesh
			system->StopBenchmark();
			system->StartBenchmark("TraceMesh");
			int gPtr = 0; int start = 0;
			while (gPtr < groups.Count())
			{
				int end = start + groups[gPtr].Count;
				while (gPtr + 1 < groups.Count() && groups[gPtr+1].GroupId == groups[gPtr].GroupId)
				{
					gPtr++;
					end += groups[gPtr].Count;
				}
				// MeshTrace skeleton[start, end]
				if (groups[gPtr].GroupId != -1)
				{
					done = false;
					meshes[groups[gPtr].GroupId]->ComputeIntersectionGPU(system, start, end, skeleton, continuation, Instances, rays, rayDiffs, intermediateIntersections, result);
				}
				start = end;
				gPtr++;
			}
			system->StopBenchmark();
		} while (!done);
	}
}